package com.hxx.controller;

import lombok.SneakyThrows;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.InputStreamResource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;

@RestController
public class QinwenController3 {

    // 注入VectorStore实例(向量数据库)
    @Autowired
    private VectorStore vectorStore;

    // 读取文件并返回文本内容
    @SneakyThrows
    @PostMapping("/embedding")
    public String readFile(MultipartFile file) {
        try {
            InputStreamResource resource = new InputStreamResource(file.getInputStream());
            // 从IO流中读取文件
            TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(resource);
            // 将文本内容划分成更小的块
            List<Document> splitDocuments = new TokenTextSplitter()
                    // 设置每个块的最大长度
                    .apply(tikaDocumentReader.read());
            // 将划分后的文本块存储在向量数据库中
            System.out.println(splitDocuments);
            vectorStore.add(splitDocuments);
            return "存储在向量数据库中";
        }catch (Exception e){
            e.printStackTrace();
        }
        return "存储失败";
    }

    /**
     * 查询向量数据库
     *
     * @param query 用户的提问
     * @return 匹配到的文档
     */

    @GetMapping("query")
    public List<Document> query(@RequestParam String query) {
        return vectorStore.similaritySearch(query);
    }
}